﻿using System;
using System.Buffers;
using System.IO;

namespace ManagedTFRecordLibrary
{
    public static class TFRecordReaderExtensions
    {

        public static Example TryParseExample(this TFRecordReader reader)
        {
            if (reader == null)
            {
                throw new ArgumentNullException(nameof(reader));
            }

            int length = (int)reader.CurrentLength;
            byte[] buffer = ArrayPool<byte>.Shared.Rent(length);
            try
            {
                if (!reader.TryReadCurrent(new Span<byte>(buffer, 0, length)))
                {
                    return null;
                }

                using (var ms = new MemoryStream(buffer,0, length))
                {
                    var tfExample = ProtoBuf.Serializer.Deserialize<Tensorflow.Example>(ms);
                    return ExampleConverter.ConvertFromTFExample(tfExample);
                }
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }

        }
    }
}
